import cv2
import numpy as np
import matplotlib.pyplot as plt

from pathlib import Path
from tensorflow.keras.optimizers import SGD
from sklearn.metrics import classification_report
from sklearn.preprocessing import OneHotEncoder
from sklearn.model_selection import train_test_split

from lenet import Lenet


def main():
    img_dir = Path("../../dataset/dog_cat")
    data = []
    labels = []
    for num, i in enumerate(img_dir.iterdir(), start=1):
        raw_data = cv2.resize(cv2.imread(str(i)), (32, 32),
                              interpolation=cv2.INTER_AREA)
        labels.append(i.name.split('.')[0])
        data.append(raw_data)
        print(f"\r共加载{num}张图片", end='', flush=True)
    data = np.array(data).astype("float") / 255.0
    labels = np.array(labels)

    train_x, test_x, train_y, test_y = train_test_split(data, labels,
                                                        test_size=0.25)
    # 编码
    oh = OneHotEncoder()
    train_y = oh.fit_transform(np.array([[i] for i in train_y])).toarray()
    test_y = oh.transform(np.array([[i] for i in test_y])).toarray()

    # 模型编译
    model = Lenet.build(width=32, height=32, depth=3, classes=2)
    opt = SGD(0.01)
    model.compile(loss="binary_crossentropy", optimizer=opt,
                  metrics=["acc"])

    # 模型训练
    print('[info]:开始训练...')
    record = model.fit(train_x, train_y, validation_data=(test_x, test_y),
                       batch_size=64, epochs=100, verbose=1)

    # 模型评估
    print('[info]:开始评估...')
    predictions = model.predict(test_x, batch_size=64, verbose=1)
    print(classification_report(test_y.argmax(1),
                                predictions.argmax(1),
                                target_names=oh.categories_[0]))

    # 画图
    plt.style.use("ggplot")
    plt.figure()
    plt.plot(np.arange(0, 100), record.history["loss"], label="train_loss")
    plt.plot(np.arange(0, 100), record.history["val_loss"], label="val_loss")
    plt.plot(np.arange(0, 100), record.history["acc"], label="train_acc")
    plt.plot(np.arange(0, 100), record.history["val_acc"], label="val_acc")
    plt.title("Training Loss and Accuracy")
    plt.xlabel("Epoch #")
    plt.ylabel("Loss/Accuracy")
    plt.legend()
    plt.show()


if __name__ == '__main__':
    main()
